import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from raw.base_model import BaseModel
from raw.models.MNIST import training_params

def get_model_group():
    return [MNIST_A(), MNIST_B(), MNIST_C(), MNIST_D()]

def get_optimizer(model):
    return optim.SGD(model.parameters(),
                    lr = training_params['learning_rate'],
                    momentum = training_params['momentum'],
                    weight_decay = training_params['decay'],
                    nesterov = True)

class MNIST_A(BaseModel):
    def __init__(self):
        super().__init__('MNIST', None)

        self.conv32 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3),
            nn.ReLU(),
            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.conv64 = nn.Sequential(
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3),
            nn.ReLU(),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.fc1 = nn.Linear(4 * 4 * 64, 200)
        self.dropout = nn.Dropout2d(p=0.5)
        self.fc2 = nn.Linear(200, 200)
        self.fc3 = nn.Linear(200, 10)

    def forward(self, x):
        out = self.conv32(x)
        out = self.conv64(out)
        out = out.view(-1, 4 * 4 * 64)
        out = F.relu(self.fc1(out))
        out = self.dropout(out)
        out = F.relu(self.fc2(out))
        out = self.fc3(out)
        return out

    def get_optimizer(self):
        return get_optimizer(self)


class MNIST_B(BaseModel):
    def __init__(self):
        super().__init__('MNIST', None)

        self.conv64 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=64, kernel_size=8, stride=(2, 2), padding=(17, 17)),
            nn.ReLU()
        )
        self.conv128_1 = nn.Sequential(
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=6, stride=(2, 2)),
            nn.ReLU()
        )
        self.conv128_2 = nn.Sequential(
            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=5),
            nn.ReLU()
        )
        self.fc = nn.Linear(8 * 8 * 128, 10)
        self.dropout_1 = nn.Dropout2d(p=0.2)
        self.dropout_2 = nn.Dropout2d(p=0.5)

    def forward(self, x):
        out = self.dropout_1(x)
        out = self.conv64(out)
        out = self.conv128_1(out)
        out = self.conv128_2(out)
        out = self.dropout_2(out)
        out = out.view(-1, 8 * 8 * 128)
        out = self.fc(out)
        return out

    def get_optimizer(self):
        return get_optimizer(self)


class MNIST_C(BaseModel):
    def __init__(self):
        super().__init__('MNIST', None)

        self.conv128 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=128, kernel_size=3),
            nn.ReLU()
        )
        self.conv64 = nn.Sequential(
            nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3),
            nn.ReLU()
        )
        self.fc1 = nn.Linear(24 * 24 * 64, 128)
        self.fc2 = nn.Linear(128, 10)
        self.dropout_1 = nn.Dropout2d(p=0.25)
        self.dropout_2 = nn.Dropout2d(p=0.5)

    def forward(self, x):
        out = self.conv128(x)
        out = self.conv64(out)
        out = self.dropout_1(out)
        out = out.view(-1, 24 * 24 * 64)
        out = F.relu(self.fc1(out))
        out = self.dropout_2(out)
        out = self.fc2(out)
        return out

    def get_optimizer(self):
        return get_optimizer(self)


class MNIST_D(BaseModel):
    def __init__(self):
        super().__init__('MNIST', None)

        self.fc1 = nn.Linear(28 * 28 * 1, 300)
        self.fc2 = nn.Linear(300, 300)
        self.fc3 = nn.Linear(300, 10)
        self.dropout = nn.Dropout2d(p=0.5)

    def forward(self, x):
        out = x.view(-1, 28 * 28 * 1)
        out = F.relu(self.fc1(out))
        out = self.dropout(out)
        out = F.relu(self.fc2(out))
        out = self.dropout(out)
        out = F.relu(self.fc2(out))
        out = self.dropout(out)
        out = F.relu(self.fc2(out))
        out = self.dropout(out)
        out = self.fc3(out)
        return out

    def get_optimizer(self):
        return get_optimizer(self)
